{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# analysis-get-calibration_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import tvm\n",
    "import tvm.relay.testing\n",
    "from tvm import relay\n",
    "from tvm.relay import transform\n",
    "from tvm.relay.analysis import get_calibration_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "{func}`get_calibration_data` 的函数，它用于获取给定 `relay` 计算图的校准数据。校准数据包括每个函数的输入和输出值。返回的数据使用每个函数的 GlobalVar 作为键。用户可以通过使用 `inputs` 或 `outputs` 作为键来进一步访问输入和输出。\n",
    "\n",
    "以下是一些限制：\n",
    "1. 输入模块（计算图）不能有控制流。\n",
    "2. 每个函数的输入参数不能是元组（输出可以是元组）。\n",
    "3. 只处理顶层（top-level）函数（即不处理嵌套函数）。\n",
    "4. 只处理具有 `Compiler` 属性设置的函数。\n",
    "\n",
    "函数接受两个参数：\n",
    "- `mod`：{class}`tvm.IRModule` 类型，用于收集校准数据的输入模块。\n",
    "- `data`：`Dict[str, NDArray]` 类型，用于运行模块的输入数据。\n",
    "\n",
    "函数返回一个字典，其键为 {class}`tvm.relay.GlobalVar`，值为包含输入和输出数据的字典。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_data_size(mod, data):\n",
    "    assert len(data) == len(mod.functions) - 1\n",
    "    for key, value in mod.functions.items():\n",
    "        if key.name_hint != \"main\":\n",
    "            assert len(data[key][\"inputs\"]) == len(value.params)\n",
    "            if isinstance(value.body, relay.Tuple):\n",
    "                assert len(data[key][\"outputs\"]) == len(value.body.fields)\n",
    "            else:\n",
    "                assert len(data[key][\"outputs\"]) == 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试简单计算图校准数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A module with two subgraphs\n",
    "mod = tvm.IRModule()\n",
    "\n",
    "x0 = relay.var(\"x0\", shape=(8, 8))\n",
    "y0 = relay.var(\"y0\", shape=(8, 8))\n",
    "z0 = x0 + y0\n",
    "z1 = x0 - y0\n",
    "z2 = relay.Tuple((z0, z1))\n",
    "f0 = relay.Function([x0, y0], z2)\n",
    "f0 = f0.with_attr(\"Compiler\", \"test_graph\")\n",
    "g0 = relay.GlobalVar(\"g0\")\n",
    "mod[g0] = f0\n",
    "mod = relay.transform.InferType()(mod)\n",
    "\n",
    "x1 = relay.var(\"x1\", shape=(8, 8))\n",
    "y1 = relay.var(\"y1\", shape=(8, 8))\n",
    "z1 = x1 - y1\n",
    "f1 = relay.Function([x1, y1], z1)\n",
    "f1 = f1.with_attr(\"Compiler\", \"test_graph\")\n",
    "g1 = relay.GlobalVar(\"g1\")\n",
    "mod[g1] = f1\n",
    "mod = relay.transform.InferType()(mod)\n",
    "\n",
    "x = relay.var(\"x\", shape=(8, 8))\n",
    "y = relay.var(\"y\", shape=(8, 8))\n",
    "z = relay.var(\"z\", shape=(8, 8))\n",
    "c0 = relay.Call(g0, [x, y])\n",
    "c1 = relay.Call(g1, [relay.TupleGetItem(c0, 0), z])\n",
    "fm = relay.Function([x, y, z], c1)\n",
    "mod[\"main\"] = fm\n",
    "mod = relay.transform.InferType()(mod)\n",
    "\n",
    "x_data = np.random.rand(8, 8).astype(\"float32\")\n",
    "y_data = np.random.rand(8, 8).astype(\"float32\")\n",
    "z_data = np.random.rand(8, 8).astype(\"float32\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@g0</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x0: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y0: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, Compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_graph&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32], Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32]) {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> add(<span style=\"color: #A2F; font-weight: bold\">%</span>x0, <span style=\"color: #A2F; font-weight: bold\">%</span>y0) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> subtract(<span style=\"color: #A2F; font-weight: bold\">%</span>x0, <span style=\"color: #A2F; font-weight: bold\">%</span>y0) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  (<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32], Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32]) <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@g1</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x1: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y1: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, Compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_graph&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] {\n",
       "  subtract(<span style=\"color: #A2F; font-weight: bold\">%</span>x1, <span style=\"color: #A2F; font-weight: bold\">%</span>y1) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>z: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F\">@g0</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span>y) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32], Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32]) <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2.0</span> <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F\">@g1</span>(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>z) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = get_calibration_data(mod, {\"x\": x_data, \"y\": y_data, \"z\": z_data})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys([I.GlobalVar(\"g0\"), I.GlobalVar(\"g1\")])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the number and orders\n",
    "check_data_size(mod, data)\n",
    "tvm.testing.assert_allclose(data[g0][\"inputs\"][0].numpy(), x_data)\n",
    "tvm.testing.assert_allclose(data[g0][\"inputs\"][1].numpy(), y_data)\n",
    "tvm.testing.assert_allclose(data[g0][\"outputs\"][0].numpy(), x_data + y_data)\n",
    "tvm.testing.assert_allclose(data[g0][\"outputs\"][1].numpy(), x_data - y_data)\n",
    "tvm.testing.assert_allclose(data[g1][\"inputs\"][0].numpy(), x_data + y_data)\n",
    "tvm.testing.assert_allclose(data[g1][\"inputs\"][1].numpy(), z_data)\n",
    "tvm.testing.assert_allclose(data[g1][\"outputs\"][0].numpy(), x_data + y_data - z_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.relay.analysis import _ffi_api"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map = _ffi_api.get_calibrate_output_map(mod)\n",
    "run_mod = _ffi_api.get_calibrate_module(mod)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{I.GlobalVar(\"g0\"): [0, 2, 2], I.GlobalVar(\"g1\"): [4, 2, 1]}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@g0</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x0: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y0: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, Compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>, Inline<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32], Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32]) {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> add(<span style=\"color: #A2F; font-weight: bold\">%</span>x0, <span style=\"color: #A2F; font-weight: bold\">%</span>y0) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> subtract(<span style=\"color: #A2F; font-weight: bold\">%</span>x0, <span style=\"color: #A2F; font-weight: bold\">%</span>y0) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  (<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32], Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32]) <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@g1</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x1: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y1: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, Compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>, Inline<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] {\n",
       "  subtract(<span style=\"color: #A2F; font-weight: bold\">%</span>x1, <span style=\"color: #A2F; font-weight: bold\">%</span>y1) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>z: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>) {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F\">@g0</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span>y) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32], Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32]) <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2.0</span> <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2.0</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2.1</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">6</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F\">@g1</span>(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>z) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  (<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span>y, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>z, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">6</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "run_mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@g0</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x0: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y0: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, Compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_graph&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> (Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32], Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32]) {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> add(<span style=\"color: #A2F; font-weight: bold\">%</span>x0, <span style=\"color: #A2F; font-weight: bold\">%</span>y0) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> subtract(<span style=\"color: #A2F; font-weight: bold\">%</span>x0, <span style=\"color: #A2F; font-weight: bold\">%</span>y0) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  (<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32], Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32]) <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@g1</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x1: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y1: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, Compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_graph&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] {\n",
       "  subtract(<span style=\"color: #A2F; font-weight: bold\">%</span>x1, <span style=\"color: #A2F; font-weight: bold\">%</span>y1) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>z: Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F\">@g0</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span>y) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32], Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32]) <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2.0</span> <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F\">@g1</span>(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>z) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试 DNNL 校准数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_mobilenet_dnnl():\n",
    "    if not tvm.get_global_func(\"relay.ext.dnnl\", True):\n",
    "        print(\"skip because DNNL codegen is not available\")\n",
    "        return\n",
    "\n",
    "    dtype = \"float32\"\n",
    "    ishape = (1, 3, 224, 224)\n",
    "    mod, params = relay.testing.mobilenet.get_workload(batch_size=1, dtype=\"float32\")\n",
    "\n",
    "    mod = transform.AnnotateTarget([\"dnnl\"])(mod)\n",
    "    mod = transform.MergeCompilerRegions()(mod)\n",
    "    mod = transform.PartitionGraph()(mod)\n",
    "\n",
    "    i_data = np.random.uniform(0, 1, ishape).astype(dtype)\n",
    "    data = get_calibration_data(mod, {\"data\": i_data, **params})\n",
    "\n",
    "    # Check the number and orders\n",
    "    check_data_size(mod, data)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `src/relay/analysis/get_calibration_data.cc`\n",
    "\n",
    "为了获取校准数据，需要进行两个步骤：\n",
    "\n",
    "1. 准备生成张量值的模块（`GetCalibrateModule`）。\n",
    "2. 生成值与函数之间的映射（`GetCalibrateOutputMap`）。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```cpp\n",
    "\n",
    "/*!\n",
    " * \\brief This function returns a module that will be used by\n",
    " * the relay graph executor for collecting the calibration data.\n",
    " * To do that, we first make all inputs and outputs of each\n",
    " * function into the final output (i.e., the final output is a\n",
    " * tuple of tensors). Then, we change the compiler attribute of\n",
    " * each function. Finally, we mark all function to be inlined.\n",
    " */\n",
    "\n",
    "class Collector : public ExprRewriter {\n",
    " public:\n",
    "  explicit Collector(const IRModule& module) : module_(module) {}\n",
    "\n",
    "  Expr Rewrite_(const CallNode* call, const Expr& post) final {\n",
    "    // check if the function implementation is available\n",
    "    // intrinsic functions are excluded for now\n",
    "    if (call->op->IsInstance<GlobalVarNode>()) {\n",
    "      auto var = Downcast<GlobalVar>(call->op);\n",
    "      ICHECK(module_->ContainGlobalVar(var->name_hint)) << \"Function \" << var << \" is not defined\";\n",
    "      // we only handle functions with Compiler attribute set\n",
    "      auto func = Downcast<Function>(module_->Lookup(var));\n",
    "      if (func->GetAttr<String>(attr::kCompiler)) {\n",
    "        // collect all the inputs and outputs\n",
    "        for (const auto& it : call->args) new_outputs_.push_back(it);\n",
    "        new_outputs_.push_back(post);\n",
    "      }\n",
    "    }\n",
    "    return post;\n",
    "  }\n",
    "\n",
    "  Array<Expr> GetNewOutputs() { return new_outputs_; }\n",
    "\n",
    " private:\n",
    "  const IRModule& module_;\n",
    "  Array<Expr> new_outputs_;\n",
    "};\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Collector` 类，它继承自 `ExprRewriter`。`Collector` 类的主要目的是收集校准数据，以便在 `relay` 计算图执行器中使用。为了实现这个目标，`Collector` 类首先将所有函数的输入和输出转换为最终输出（即一个张量元组）。然后，它更改每个函数的编译器属性。最后，它将所有函数标记为内联。\n",
    "\n",
    "`Collector` 类有一个构造函数，接受 `IRModule` 类型的参数。它还重写了 `ExprRewriter` 类的 `Rewrite_` 方法，该方法用于处理 `CallNode` 类型的表达式。在 `Rewrite_` 方法中，首先检查函数实现是否可用，排除了 `intrinsic` 函数。然后，只处理设置了 `Compiler` 属性的函数。对于这些函数，它会收集所有的输入和输出，并将它们添加到 `new_outputs_` 数组中。最后，`GetNewOutputs` 方法返回收集到的新输出数组。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```cpp\n",
    "\n",
    "Expr FlattenOutputTuple(const Array<Expr>& exprs) {\n",
    "  Array<Expr> fields;\n",
    "  for (const auto& it : exprs) {\n",
    "    ICHECK(it->checked_type_.defined());\n",
    "    if (auto* tn = it->checked_type_.as<TupleTypeNode>()) {\n",
    "      // TODO(seanlatias): for now input argument cannot be a tuple\n",
    "      ICHECK(it->IsInstance<CallNode>());\n",
    "      for (size_t i = 0; i < tn->fields.size(); i++) {\n",
    "        fields.push_back(TupleGetItem(it, i));\n",
    "      }\n",
    "    } else {\n",
    "      fields.push_back(it);\n",
    "    }\n",
    "  }\n",
    "  return Tuple(fields);\n",
    "}\n",
    "\n",
    "IRModule GetCalibrateModule(IRModule module) {\n",
    "  auto glob_funcs = module->functions;\n",
    "  // module is mutable, hence, we make a copy of it.\n",
    "  module.CopyOnWrite();\n",
    "  for (const auto& pair : glob_funcs) {\n",
    "    if (auto opt = pair.second.as<Function>()) {\n",
    "      // we only collect the outputs for main function\n",
    "      if (pair.first->name_hint == \"main\") {\n",
    "        auto func = opt.value();\n",
    "        Collector collector(module);\n",
    "        PostOrderRewrite(func->body, &collector);\n",
    "        auto new_outputs = collector.GetNewOutputs();\n",
    "        Expr tuple = FlattenOutputTuple(new_outputs);\n",
    "        func = Function(func->params, tuple, tuple->checked_type_, func->type_params, func->attrs);\n",
    "        module->Update(pair.first, func);\n",
    "      }\n",
    "    }\n",
    "  }\n",
    "  // reset the attribute of functions for running graph executor\n",
    "  for (const auto& pair : glob_funcs) {\n",
    "    if (auto opt = pair.second.as<Function>()) {\n",
    "      auto func = opt.value();\n",
    "      if (func->GetAttr<String>(attr::kCompiler)) {\n",
    "        // we need to inline the functions in order to run grpah runtime\n",
    "        func = WithAttr(std::move(func), attr::kInline, tvm::Integer(1));\n",
    "        // reset the compiler attribute to null for llvm execution\n",
    "        func = WithAttr(std::move(func), attr::kCompiler, NullValue<ObjectRef>());\n",
    "        module->Update(pair.first, func);\n",
    "      }\n",
    "    }\n",
    "  }\n",
    "  return module;\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这段代码定义了两个函数：`FlattenOutputTuple` 和 `GetCalibrateModule`。\n",
    "\n",
    "`FlattenOutputTuple` 函数接受 `Expr` 类型的数组作为参数，然后遍历这个数组。对于每个元素，它首先检查元素的类型是否已定义。如果元素是一个元组类型，它会将元组中的每个元素添加到一个新的 `fields` 数组中。如果元素不是元组类型，它会直接将元素添加到 `fields` 数组中。最后，函数返回一个包含 `fields` 数组所有元素的新元组。\n",
    "\n",
    "`GetCalibrateModule` 函数接受 `IRModule` 类型的参数，并对其进行处理以生成校准模块。首先，它获取模块中的所有全局函数。然后，它创建可写的副本，以便在不改变原始模块的情况下进行修改。接下来，它遍历所有全局函数，只处理名为 `\"main\"` 的函数。对于每个符合条件的函数，它使用 `Collector` 类收集函数体中的所有输出，并将这些输出展平为一个元组。然后，它用新的输出元组替换原来的函数体，并将更新后的函数放回模块中。最后，它重置所有具有编译器属性的函数的属性，以便在运行图执行器时可以内联这些函数，并将编译器属性设置为 `null` 以便于 LLVM 执行。函数返回处理后的模块。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```cpp\n",
    "/*!\n",
    " * \\brief This function generates the output mapping between\n",
    " * the calibration data and each function. The key is a\n",
    " * GlobalVar that corresponds to each function and the value\n",
    " * is an array of integers. The size of the array is always\n",
    " * three. The first value is the offset the points to the start.\n",
    " * The second value is the number of inputs. The third value\n",
    " * is the number of outputs.\n",
    " */\n",
    "\n",
    "class OutputMapper : public ExprRewriter {\n",
    " public:\n",
    "  OutputMapper(Map<GlobalVar, Array<Integer>>* output_map, const IRModule& module, size_t* offset)\n",
    "      : output_map_(output_map), module_(module), offset_(offset) {}\n",
    "\n",
    "  Expr Rewrite_(const CallNode* call, const Expr& post) final {\n",
    "    if (call->op->IsInstance<GlobalVarNode>()) {\n",
    "      auto var = Downcast<GlobalVar>(call->op);\n",
    "      ICHECK(module_->ContainGlobalVar(var->name_hint)) << \"Function \" << var << \" is not defined\";\n",
    "      ICHECK_EQ(output_map_->count(var), 0)\n",
    "          << \"Repeated function call \" << var << \" is not supported.\";\n",
    "      auto func = Downcast<Function>(module_->Lookup(var));\n",
    "      // we only handle functions with Compiler attribute set\n",
    "      if (func->GetAttr<String>(attr::kCompiler)) {\n",
    "        Array<Integer> info;\n",
    "        // the first value is the offset\n",
    "        info.push_back(Integer(*offset_));\n",
    "        // the second value is the number of inputs\n",
    "        info.push_back(Integer(call->args.size()));\n",
    "        // the third value is the number of outputs\n",
    "        // we need to check if the output is a tuple\n",
    "        size_t out_size = 1;\n",
    "        if (auto* tn = func->body.as<TupleNode>()) {\n",
    "          info.push_back(Integer(tn->fields.size()));\n",
    "          out_size = tn->fields.size();\n",
    "        } else {\n",
    "          info.push_back(Integer(1));\n",
    "        }\n",
    "        output_map_->Set(var, info);\n",
    "        // calculate the offset for the next function\n",
    "        *offset_ = *offset_ + call->args.size() + out_size;\n",
    "      }\n",
    "    }\n",
    "    return post;\n",
    "  }\n",
    "\n",
    " private:\n",
    "  Map<GlobalVar, Array<Integer>>* output_map_;\n",
    "  const IRModule& module_;\n",
    "  size_t* offset_;\n",
    "};\n",
    "\n",
    "Map<GlobalVar, Array<Integer>> GetCalibrateOutputMap(const IRModule& module) {\n",
    "  Map<GlobalVar, Array<Integer>> output_map;\n",
    "  size_t offset = 0;\n",
    "  auto glob_funcs = module->functions;\n",
    "  for (const auto& pair : glob_funcs) {\n",
    "    if (const auto* func = pair.second.as<FunctionNode>()) {\n",
    "      if (pair.first->name_hint == \"main\") {\n",
    "        OutputMapper output_mapper(&output_map, module, &offset);\n",
    "        PostOrderRewrite(func->body, &output_mapper);\n",
    "      }\n",
    "    }\n",
    "  }\n",
    "\n",
    "  return output_map;\n",
    "}\n",
    "\n",
    "TVM_REGISTER_GLOBAL(\"relay.analysis.get_calibrate_module\").set_body_typed([](IRModule mod) {\n",
    "  return GetCalibrateModule(mod);\n",
    "});\n",
    "\n",
    "TVM_REGISTER_GLOBAL(\"relay.analysis.get_calibrate_output_map\")\n",
    "    .set_body_typed([](const IRModule& mod) { return GetCalibrateOutputMap(mod); });\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`OutputMapper` 类，它继承自 `ExprRewriter`。`OutputMapper` 的主要功能是生成校准数据与每个函数之间的输出映射。映射的键是 `GlobalVar`，对应于每个函数；值是一个整数数组，大小始终为 3。第一个值是指向起始位置的偏移量；第二个值是输入的数量；第三个值是输出的数量。\n",
    "\n",
    "`GetCalibrateOutputMap` 函数接受 `IRModule` 类型的参数，并返回一个包含校准数据的映射。它首先创建一个空的 `output_map` 和一个初始偏移量 `offset`。然后，遍历模块中的所有全局函数，只处理名为 `\"main\"` 的函数。对于每个符合条件的函数，使用 `OutputMapper` 类收集函数体中的所有输出，并将这些输出展平为一个元组。最后，将更新后的函数放回模块中。\n",
    "\n",
    "全局函数：`relay.analysis.get_calibrate_module`和`relay.analysis.get_calibrate_output_map` 分别用于获取校准模块和校准输出映射。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aix",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
